﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Athena;
using Athena.Plugins;
using Athena.Core;
using Athena.Functions;
using Troschuetz.Random;

namespace StochasticSimulator
{
    public class Simulator : Athena.Plugins.Simulator, Athena.Plugins.Analyzer
    {
        protected static MT19937Generator rUnif = new MT19937Generator();

        Module targetModule = null;
        double startTime = 0, endTime = 0, Time = 0;
        int numPoints = 1000;
        double[,] _data = null;

        public string Name
        {
            get { return "Stochastic"; }
        }

        public Module CurrentModule
        {
            get { return targetModule; }
            set
            {
                targetModule = value;
                setupMatrices();
                _selectionList = null;
                _data = null;
            }
        }
        public void loadSBML(string sSBML)
        {
        }
        public bool IsLoaded(Module module)
        {
            return targetModule == module;
        }
        public void setNumPoints(int numPoints)
        {
            this.numPoints = numPoints;
            _data = null;
        }
        public void setTimeEnd(double endTime)
        {
            this.endTime = endTime;
            _data = null;
        }
        public void setTimeStart(double startTime)
        {
            this.startTime = startTime;
            _data = null;
        }

        double[] SpeciesConc;
        public void reset()
        {
            Time = startTime;
            SpeciesConc = speciesConcentrations();
            _data = null;
        }

        public bool HasModel
        {
            get { return targetModule != null; }
        }

        public System.Collections.ArrayList getAvailableSymbols()
        {
            ArrayList oResult = new ArrayList();
            oResult.Add(new ArrayList(new object[] { "Time", new ArrayList(new object[] { "time" }) }));
            
            if (targetModule != null)
            {
                Hashtable hash = targetModule.AllSpecies;
                Hashtable hash2 = targetModule.GlobalParameters;
                Hashtable hash3 = targetModule.AllReactions;
                ArrayList floatingSpecies = new ArrayList();
                ArrayList boundarySpecies = new ArrayList();
                ArrayList parameters = new ArrayList();
                ArrayList fluxes = new ArrayList();

                //floatingSpecies.Add("Time");

                foreach (string s in hash.Keys)
                {
                    Species p = hash[s] as Species;
                    if (p != null)
                        if (p.Boundary)
                            boundarySpecies.Add(s);
                        else
                            floatingSpecies.Add(s);
                }

                foreach (string s in hash2.Keys)
                {
                    parameters.Add(s);
                }

                foreach (string s in hash3.Keys)
                {
                    fluxes.Add(s);
                }

                oResult.Add(new ArrayList ( new object[] { "Floating Species", floatingSpecies }));

                oResult.Add(new ArrayList(new object[] { "Boundary Species", boundarySpecies }));

                oResult.Add(new ArrayList(new object[] { "Fluxes", fluxes }));

                oResult.Add(new ArrayList(new object[] { "Parameters", parameters }));

            }
            return oResult;
        }
        public System.Collections.ArrayList getBoundarySpeciesNames()
        {
            if (targetModule != null)
            {
                Hashtable hash = targetModule.AllSpecies;
                ArrayList array = new ArrayList(hash.Count);
                
                foreach (string s in hash.Keys)
                {
                    Species p = hash[s] as Species;
                    if (p != null && p.Boundary)
                        array.Add(s);
                }
                return array;
            }
            return null;
        }
        public System.Collections.ArrayList getFloatingSpeciesNames()
        {
            if (targetModule != null)
            {
                Hashtable hash = targetModule.AllSpecies;
                ArrayList array = new ArrayList(hash.Count);

                foreach (string s in hash.Keys)
                {
                    Species p = hash[s] as Species;
                    if (p != null && !p.Boundary)
                        array.Add(s);
                }
                return array;
            }
            return null;
        }
        public System.Collections.ArrayList getGlobalParameterNames()
        {
            if (targetModule != null)
            {
                Hashtable hash = targetModule.GlobalParameters;
                ArrayList names = new ArrayList(hash.Count);
                
                foreach (string s in hash.Keys)
                {
                    names.Add(s);
                }
                return names;
            }
            return null;
        }
        public double[] getBoundarySpeciesConcentration()
        {
            if (targetModule != null)
            {
                Hashtable hash = targetModule.AllSpecies;
                double[] values = new double[hash.Count];
                int i = 0;
                foreach (string s in hash.Keys)
                {
                    Species p = hash[s] as Species;
                    if (p != null && p.Boundary)
                        values[i] = p.Concentration;
                    ++i;
                }
                return values;
            }
            return null;
        }
        public double[] getGlobalParameterValues()
        {
            if (targetModule != null)
            {
                Hashtable hash = targetModule.GlobalParameters;
                double[] parameters = new double[hash.Count];
                int i =0;
                foreach (string s in hash.Keys)
                {
                    Parameter p = hash[s] as Parameter;
                    if (p != null)
                        parameters[i] = p.Value;
                    ++i;
                }
                return parameters;
            }
            return null;
        }

        public double getValue(string name)
        {
            if (targetModule != null)
            {
                foreach (TAlgExpression parser in rateExpressions)
                {
                    try
                    {
                        double d = parser.getVar(name);
                        return d;
                    }
                    catch (Exception) { }
                }
            }
            return 0;
        }
        public void setValue(string name, double value)
        {
            if (targetModule != null)
            {
                foreach (TAlgExpression parser in rateExpressions)
                    parser.setVar(name, value);
                _data = null;
            }
        }

        ArrayList _selectionList = null;
        public ArrayList getSelectionList()
        {
            if (_selectionList == null)
            {
                if (N != null && N.c > 0 && N.r > 0)
                {
                    _selectionList = new ArrayList();
                    _selectionList.Add("time");
                    _selectionList.AddRange(N.rowNames);
                    _selectionList.AddRange(N.columnNames);
                    return _selectionList;
                }
            }
            return _selectionList;
        }
        public void setSelectionList(ArrayList list)
        {
            if (list != null && list.Count > 0)
                _selectionList = list;
        }
        public string SBML
        {
            get 
            {
                if (targetModule != null)
                    return Utility.GetSBML(targetModule.JarnacScript);
                else
                    return "";
            }
            set { }
        }

        /// <summary>
        /// Gillespie
        /// </summary>
        /// <returns></returns>
        public double[,] simulate()
        {
            if (N == null) return null;
            zeroPropensity = false;

            int numSpecies = N.rows, numReactions = N.cols;
            SpeciesConc = speciesConcentrations();
            Time = startTime;
            
            double[,] results = _data;
            double[,] results2;
            
            int iter = 0, i = 0;

            if (startTime > endTime)
            {
                double temp = startTime;
                startTime = endTime;
                endTime = temp;
            }
            if (numPoints < endTime - startTime) numPoints = (int)(endTime - startTime);
            double dt = (endTime - startTime) / numPoints;
            
            if (results == null)
            {
                results = new double[numPoints, numSpecies + numReactions + 1];

                double[] rates = null;
                double lambda;

                for (i = 0; i < SpeciesConc.Length; ++i) SpeciesConc[i] = Math.Ceiling(SpeciesConc[i]);

                while (Time < endTime)
                {
                    if (iter >= numPoints)
                    {
                        numPoints *= 2;
                        results2 = results;
                        results = new double[numPoints, results.GetLength(1)];
                        for (i = 0; i < iter; ++i)
                            for (int j = 0; j < results.GetLength(1); ++j)
                                results[i, j] = results2[i, j];
                    }

                    if (TimeNodes != null)
                        for (i = 0; i < TimeNodes.Length; ++i)
                            TimeNodes[i].value = Time;
                     
                    rates = getRates(SpeciesConc);

                    if (Time > iter * dt)  //save data
                    {
                        results[iter, 0] = Time;
                        for (i = 0; i < numSpecies; ++i)
                            results[iter, 1 + i] = SpeciesConc[i];

                        for (i = 0; i < rates.Length; ++i)
                            results[iter, 1 + numSpecies + i] = rates[i]; //save rates

                        ++iter;
                    }

                    lambda = 0;
                    for (i = 0; i < rates.Length; ++i)
                    {
                        lambda += rates[i];  //get raw rates                        
                    }

                    if (lambda == 0)         //no more species
                    {
                        Time += dt;            //maybe something will happen later...
                    }
                    else
                    {

                        for (i = 0; i < rates.Length; ++i) rates[i] /= lambda;  //convert to probabilities 

                        Time += -Math.Log(rUnif.NextDouble()) / lambda; //poisson time increment

                        double rand = rUnif.NextDouble();  //spin the wheel
                        double total = 0;

                        for (i = 0; i < rates.Length; ++i)  //see where it stops
                        {
                            if (rand > total && rand < (total + rates[i])) //reaction #k
                                break;
                            total += rates[i];
                        }
                        if (i < rates.Length) //perform that reaction 
                        {
                            for (int j = 0; j < numSpecies; ++j)
                            {
                                SpeciesConc[j] += N[j, i];
                            }
                        }
                    }
                }
            }
            else
            {
                iter = _data.GetLength(0);
            }

            if (_data != results)
            {
                int n = results.GetLength(1);
                _data = new double[iter, n];
                for (i = 0; i < iter; ++i)
                    for (int j = 0; j < n; ++j)
                        _data[i, j] = results[i, j];
            }

            if (_selectionList != null && _selectionList.Count > 1 && results.GetLength(1) == 1 + N.cols + N.rows)
            {
                ArrayList newSelectionList = new ArrayList(_selectionList.Count);

                if (_selectionList.Contains("time") || _selectionList.Contains("Time"))
                    newSelectionList.Add("time");

                for (int j = 0; j < results.GetLength(1) - 1; ++j)
                {
                    if (j < N.rows && _selectionList.Contains(N.rowNames[j]))
                    {
                        newSelectionList.Add(N.rowNames[j]);
                    }
                    else
                        if (j >= N.rows && j - N.rows < N.cols && _selectionList.Contains(N.columnNames[j - N.rows]))
                        {
                            newSelectionList.Add(N.columnNames[j - N.rows]);
                        }
                }


                results2 = new double[iter, newSelectionList.Count];
                for (i = 0; i < iter; ++i)
                {
                    int k = 0;
                    if (newSelectionList.Contains("time") || newSelectionList.Contains("Time"))
                    {
                        results2[i, 0] = results[i, 0];
                        ++k;
                    }

                    for (int j = 0; j < results.GetLength(1) - 1; ++j)
                    {
                        if (j < N.rows && newSelectionList.Contains(N.rowNames[j]))
                        {
                            results2[i, k] = results[i, j + 1]; //+1 for time
                            ++k;
                        }
                        else
                            if (j >= N.rows && j - N.rows < N.cols && newSelectionList.Contains(N.columnNames[j - N.rows]))
                            {
                                results2[i, k] = results[i, j + 1]; //+1 for time
                                ++k;
                            }
                    }
                }
                _selectionList = newSelectionList;
            }
            else
            {
                results2 = new double[iter, results.GetLength(1)];
                for (i = 0; i < iter; ++i)
                    for (int j = 0; j < results.GetLength(1); ++j)
                        results2[i, j] = results[i, j];
            }

            return results2;
        }

        public bool zeroPropensity = false;
        double[] getRates(double[] speciesConc)
        {
            double[] rates = new double[N.cols];

            for (int i = 0; i < N.cols; ++i)
            {
                for (int j = 0; j < N.rows; ++j)  //set values
                {
                    if (speciesNodes[j, i] != null)
                        rateExpressions[i].setValue(speciesNodes[j, i], speciesConc[j]);
                }
                rates[i] = rateExpressions[i].eval();
                if (rates[i] < 0)
                {
                    rates[i] = 0.0d; //force non-negative
                    zeroPropensity = true; //error
                }
            }
            return rates;
        }

        double[] speciesConcentrations()
        {
            if (N == null || targetModule == null) return null;
            double[] speciesConc = new double[N.rows];
            Hashtable hash = targetModule.AllSpecies;
            for (int i = 0; i < N.r; ++i)
            {
                string s = N.rowNames[i];
                Species p = hash[s] as Species;
                if (p != null)
                {
                    speciesConc[i] = p.Concentration;
                }
            }
            return speciesConc;
        }

        TTreeNode[,] speciesNodes;        
        TAlgExpression[] rateExpressions;
        TTreeNode[] TimeNodes;
        Matrix N;

        /// <summary>
        /// assign value to each parameter in the expression tree
        /// by looking at the GlobalParameters of the module
        /// </summary>
        /// <param name="expression"></param>
        void setupParameters(TAlgExpression expression, Hashtable parameters, Hashtable species)
        {
            if (parameters == null || expression == null) return;

            TTreeNode[] symbols = expression.GetAllSymbols();

            foreach (TTreeNode tn in symbols)
            {
                string s = tn.name;
                if (parameters.ContainsKey(s) && parameters[s] is Parameter)
                {
                    Parameter p = (Parameter)parameters[s];
                    expression.setVar(s, p.Value);
                }
                if (species.ContainsKey(s) && species[s] is Species)  //initialize boundary
                {
                    Species p = (Species)species[s];
                    expression.setVar(s, p.Concentration);
                }
            }
        }

        string insertGaps(string rate)
        {
            string rate2 = ' ' + rate + ' ';
            string[] ops = TScanner.Operators;
            foreach (string op in ops)
            {
                rate2 = rate2.Replace(op, ' ' + op + ' ');
            }
            return rate2;
        }

        string replacePartWithPoPS(string rate, Hashtable parts)
        {
            foreach (string s in parts.Keys)
            {
                Part p = parts[s] as Part;
                if (p != null)
                {
                    if (rate.Contains(' ' + s + ' '))
                    {
                        string pops = insertGaps(p.PoPS.GetRate(targetModule));                        

                        if (!pops.Contains(' ' + s + ' '))
                        {
                            rate = rate.Replace(' ' + s + ' ', pops);
                            return replacePartWithPoPS(rate, parts);  //restart
                        }
                    }
                }
            }
            return rate;
        }

        string replaceRules(string rate, Hashtable rules)
        {
            foreach (string s in rules.Keys)
            {
                string p = rules[s] as string;
                if (p != null)
                {
                    if (rate.Contains(' ' + s + ' '))
                    {
                        if (!p.Contains(' ' + s + ' '))
                        {
                            rate = rate.Replace(' ' + s + ' ', "(" + p + ")");
                            return replaceRules(rate, rules);  //restart
                        }
                    }
                }
            }
            return rate;
        }

        public ArrayList getReactionNames()
        {
            ArrayList names = new ArrayList();
            if (N != null)
            {
                names.AddRange(N.columnNames);
            }
            return names;
        }

        public double[,] getRatePlots(double start, double end, int num, string variable)
        {
            if (targetModule == null) return null;

            if (start > end)
            {
                double tmp = start;
                start = end;
                end = tmp;
            }

            //setupMatrices();
            if (N == null || N.r == 0 || N.c == 0) return null;

            int index = -1;

            for (int i = 0; i < N.rowNames.Length; ++i)  //get index of species
            {
                if (N.rowNames[i] == variable)
                    index = i;
            }

            bool isParam = (index < 0);  //parameter plot

            double[,] ratePlots = new double[num, N.cols + 1];

            double[] speciesConc = speciesConcentrations();

            for (int i = 0; i < speciesConc.Length; ++i)
            {
                if (speciesConc[i] < 1) speciesConc[i] = 2;
            }
            if (Time < 1) Time = 2;

            double inc = (end - start) / num;
            double value = start;

            if (isParam)
            {
                try
                {
                    foreach (TAlgExpression parser in rateExpressions)
                    {
                        parser.setVar(variable, value);
                    }
                }
                catch (Exception)
                {
                    return null;  //no such parameter
                }
            }
            else
            {
                speciesConc[index] = value;
            }

            for (int i = 0; i < num; ++i)
            {
                double[] rates = getRates(speciesConc);

                ratePlots[i, 0] = value;

                for (int j = 0; j < N.cols && j < rates.Length; ++j)
                    ratePlots[i, 1+j] = rates[j];

                value += inc;

                if (isParam)
                {
                    foreach (TAlgExpression parser in rateExpressions)
                    {
                        parser.setVar(variable, value);
                    }
                }
                else
                {
                    speciesConc[index] = value;
                }
            }
            return ratePlots;
        }

        public double[,] getPhasePortrait(string s1, string s2, double start1, double end1, double start2, double end2, double increments)
        {
            if (targetModule == null) return null;

            int sz1 = (int)Math.Ceiling((end1 - start1) / increments),
                sz2 = (int)Math.Ceiling((end2 - start2) / increments);
            double[,] matrix = new double[sz1, sz2];

            setupMatrices();

            double[] speciesConc = speciesConcentrations();

            int index1 = -1, index2 = -1;
            for (int i = 0; i < N.rowNames.Length; ++i)  //get index of species
            {
                if (N.rowNames[i] == s1)
                    index1 = i;
                if (N.rowNames[i] == s2)
                    index2 = i;
            }

            if (index1 < 0 || index2 < 0) return null;

            double conc1 = start1, conc2 = start2;

            for (int i = 0; i < sz1; ++i) //for each conc of species1
            {
                conc1 += increments;
                conc2 = start2;

                for (int j = 0; j < sz2; ++j) //for each conc of species2
                {
                    speciesConc[index1] = conc1;
                    speciesConc[index2] = conc2;

                    double[] rates = getRates(speciesConc);
                    double d1 = 0, d2 = 0;

                    for (int n = 0; n < rates.Length; ++n)
                    {
                        d1 += N[index1, n] * rates[n];
                        d2 += N[index2, n] * rates[n];
                    }

                    if (d1 != 0)
                        matrix[i, j] = d2 / d1;
                    else
                        matrix[i, j] = 1000.0d;

                    conc2 += increments;
                }
            }

            return matrix;
        }

        /// <summary>
        /// Setup nodes and parameters
        /// </summary>
        void setupMatrices()
        {
            try
            {
                if (targetModule == null) return;
                N = targetModule.StoichiometryMatrix;
                rateExpressions = new TAlgExpression[N.cols];
                speciesNodes = new TTreeNode[N.rows, N.cols];

                List<TTreeNode> timeNodes = new List<TTreeNode>();

                Hashtable reactions = targetModule.AllReactions;
                Hashtable parameters = targetModule.GlobalParameters;
                Hashtable parts = targetModule.AllParts;
                Hashtable rules = targetModule.GlobalRules;
                Hashtable species = targetModule.AllSpecies;

                if (N.cols != reactions.Count)
                {
                    N = null;
                    return;
                }

                for (int i = 0; i < N.cols; ++i) //for each reaction
                {
                    string s = N.columnNames[i];         //reaction name
                    Edge edge = reactions[s] as Edge;   //get reaction
                    if (edge != null)
                    {
                        string rate = insertGaps(edge.GetRate(targetModule));
                        
                        rate = replacePartWithPoPS(rate, parts);
                        rate = replaceRules(rate, rules);
                        rate = rate.Replace(" ", "");

                        rateExpressions[i] = new TAlgExpression();
                        rateExpressions[i].AssignBehaviour = TAssignBehaviour.eInstallSymbols;
                        rateExpressions[i].compile(rate);  //make expression
                        setupParameters(rateExpressions[i], parameters, species);       //set param values

                        TTreeNode time = rateExpressions[i].getSymbolHandle("Time");
                        if (time != null)   //add the time variable
                            timeNodes.Add(time);
                        time = rateExpressions[i].getSymbolHandle("time");
                        if (time != null)    //add the time variable
                            timeNodes.Add(time);

                        TTreeNode[] symbols = rateExpressions[i].GetAllSymbols();

                        for (int j = 0; j < N.rows; ++j)  //for each species
                        {
                            speciesNodes[j, i] = null;

                            string s2 = N.rowNames[j];    //species name
                            foreach (TTreeNode tn in symbols) //is this species used in this rate?
                            {
                                if (s2 == tn.name)    //yes, it is
                                {
                                    TTreeNode node = rateExpressions[i].getSymbolHandle(s2);
                                    speciesNodes[j, i] = node;  //store it
                                    break;
                                }
                            }
                        }
                    }
                    else
                    {
                        N = null;
                        return;
                    }

                }
                TimeNodes = timeNodes.ToArray();
            }
            catch (Exception)
            {
            }            
        }
    }
}
